#include "stdafx.h"
#include "Common.h"
#include "WMP.h"
#include "Patch.h"
#include "WMPHook.h"

#pragma unmanaged

class WMPHookUnit {
public:
	WMPHookUnit();
	~WMPHookUnit();
};
WMPHookUnit unit;
CriticalSectionLocker* monitor;
bool Initialized = false;
bool Uninitialized = false;

RestoreModuleList* RestoreModules = NULL;
RestoreProcList* RestoreProcs = NULL;
RWLocker* Locker = NULL;
CriticalSectionLocker* CSLocker;
volatile bool HookUninitializing = false;

typedef void (WINAPI *PWMPInitialize)(HMODULE handle);
typedef void (WINAPI *PWMPInitialize2)(HMODULE module, int majorVersion, int minorVersion, IMPORT_LIBRARY lib);
typedef void (WINAPI *PWMPInitialize3)(HMODULE module, HMODULE instance, int majorVersion, int minorVersion, IMPORT_LIBRARY lib);
typedef void (WINAPI *PWMPUninitialize)();
typedef HRESULT (WINAPI *PStartThePlayer)(void*, INT_PTR, int);
typedef HRESULT (WINAPI *PCheckThePlayer)(LPSTR*, PDWORD, PDWORD, PDWORD, PDWORD);
typedef void (WINAPI *PReleaseDataObj)(LPCWSTR);
typedef HRESULT (WINAPI *PUpdateSharedFolders)();
typedef HRESULT (WINAPI *PCallFunctionElevated)(LPCWSTR);
typedef bool (WINAPI *PShouldActiveSetupBeRun)();
typedef HRESULT (WINAPI *PDllCanUnloadNow)();
typedef HRESULT (WINAPI *PDllGetClassObject)(REFCLSID rclsid, REFIID riid, LPVOID *ppv);
typedef HRESULT (WINAPI *PDllRegisterServer)();
typedef HRESULT (WINAPI *PDllUnregisterServer)();
typedef HRESULT (WINAPI *PInitializeTests)();


void ReadMethods(void* procs[], HMODULE module, const char* methods[], const int ordinals[], int count) {
	for (int k = count-1; k >= 0; k--) {
		void* proc;
		if (!methods[k])
			proc = GetProcAddress(HINST, (LPCSTR)(INT_PTR)ordinals[k]);
		else
			proc = GetProcAddress(HINST, methods[k]);
		for (int j = count-1; j > k; j--)
			if (procs[j] == proc) {
				proc = NULL;
				break;
			}
		procs[k] = proc;
	}
}

inline void FixPath(LPTSTR path) {
	size_t len = _tcslen(path);
	if (len >= 2 && path[len-1] == '\\')
		path[len-1] = 0;
}

bool LockModuleIntoProcess(HMODULE module) {
	CreatePathBuffer(pathBuffer);

	GetModuleFileName(module, pathBuffer, pathBufferLength);
#ifdef UNICODE
	void* proc = GetProcAddress(GetModuleHandle(KERNEL32DLL), "GetModuleHandleExW");
	if (proc) {
		return !!(*(PGetModuleHandleExW)proc)(GET_MODULE_HANDLE_EX_FLAG_PIN, pathBuffer, &module);
#else
	void* proc = GetProcAddress(GetModuleHandle(KERNEL32DLL), "GetModuleHandleExA");
	if (proc) {
		return !!(*(PGetModuleHandleExA)proc)(GET_MODULE_HANDLE_EX_FLAG_PIN, pathBuffer, &module);
#endif // !UNICODE
	} else {
		return !!LoadLibrary(pathBuffer);
	}
}

struct HookData {
	void* hook;
	HMODULE module;
	void* save;
};
typedef CArray<HookData> HookList;
typedef HookList* PHookList;

bool SaveMethod(LPCSTR method, void* proc, LPCSTR methods[], void* procs[], bool found, bool empty, int count) {
	if (found && !empty) return found;
	for(int k = 0; k < count; k++) {
		if (lstrcmpA(methods[k], method) == 0) {
			procs[k] = proc;
			found = true;
			if (!empty) return found;
		}
	}
	return found;
}

void SaveHookMethod(LPCSTR method, void* proc) {
	bool empty = !method && !proc;
	bool found = false;
	found = SaveMethod(method, proc, Kernel32Methods, Kernel32Hook, found, empty, Kernel32Count);
	found = SaveMethod(method, proc, user32Methods, user32Hook, found, empty, user32Count);
	found = SaveMethod(method, proc, shell32Methods, shell32Hook, found, empty, shell32Count);
	found = SaveMethod(method, proc, wmvcoreMethods, wmvcoreHook, found, empty, wmvcoreCount);
	found = SaveMethod(method, proc, msvcrtMethods, msvcrtHook, found, empty, msvcrtCount);
	found = SaveMethod(method, proc, propsysMethods, propsysHook, found, empty, propsysCount);
	found = SaveMethod(method, proc, ole32Methods, ole32Hook, found, empty, ole32Count);
	found = SaveMethod(method, proc, mfMethods, mfHook, found, empty, mfCount);
	found = SaveMethod(method, proc, quartzExternals, quartzHook, found, empty, classObjectCount);
}

void ReadHookMethods(LPCSTR methods[], void* save[], int count, List<LPCSTR, PHookList>* list) {
	char name[MAX_PATH];

	for(int k = 0; k < count; k++) {
		PHookList list2 = NULL;
		INT_PTR l = list->Add(methods[k], list2);
		for(int m = 0; m < RestoreModules->GetCount(); m++) {
			if (!RestoreModules->GetData(m))
				continue;
			HMODULE module = RestoreModules->Get(m);
			strcpy(name, "WMP");
			strcat(name, methods[k]);
			void* proc = GetProcAddress(module, name);
			if (!proc)
				continue;
			if (!list2) {
				list2 = new HookList();
				list->SetData(l, list2);
			}
			HookData hd = {proc, module, save[k]};
			list2->Add(hd);
		}
	}
}

void InitializeHookModule(HMODULE hook, HMODULE module, IMPORT_LIBRARY lib) {
	void* init = GetProcAddress(hook, "WMPInitialize3");
	if (init) {
		(*(PWMPInitialize3)init)(module, HInstance, WMPHOOK_MAJORVERSION, WMPHOOK_MINORVERSION, lib);
	} else if (init = GetProcAddress(hook, "WMPInitialize2")) {
		(*(PWMPInitialize2)init)(module, WMPHOOK_MAJORVERSION, WMPHOOK_MINORVERSION, lib);
	} else if (init = GetProcAddress(hook, "WMPInitialize")) {
		(*(PWMPInitialize)init)(module);
	}
}

void InitializeModule(HMODULE module, IMPORT_LIBRARY lib) {
	HKEY key;
	CreatePathBuffer(valueBuffer);
	CreatePathBuffer(keyBuffer);
	CreateTextBuffer(textBuffer);

	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\BM-productions\\WMPHook\\HookDlls"), 0, KEY_QUERY_VALUE, &key) == ERROR_SUCCESS) {
		try {
			DWORD count;
			if (RegQueryInfoKey(key, NULL, NULL, NULL, NULL, NULL, NULL, &count, NULL, NULL, NULL, NULL) == ERROR_SUCCESS) {
				for(DWORD k = 0; k < count; k++) {
					DWORD namelen = keyBufferLength;
					DWORD datatype;
					DWORD datalen = (valueBufferLength-1) * sizeof(TCHAR);
					valueBuffer[valueBufferLength-1] = 0;
					RegEnumValue(key, k, keyBuffer, &namelen, NULL, &datatype, (LPBYTE)valueBuffer, &datalen);
					_sntprintf(textBuffer, textBufferLength, _T("Loading hook DLL: %s\r\n"), valueBuffer);
					OutputDebugString(textBuffer);
					HMODULE hook = LoadLibrary(valueBuffer);
					PSetPrevAddress restore = NULL;
					if (!hook) continue;
					try {
						InitializeHookModule(hook, module, lib);
						restore = (PSetPrevAddress)GetProcAddress(hook, "WMPSetPrevAddress");
					} finally (
						RestoreModules->Add(hook, restore);
					)
				}
			}
		} finally (
			RegCloseKey(key);
		)
	}
	List<LPCSTR, PHookList>* list = new List<LPCSTR, PHookList>();
	try {
		ReadHookMethods(Kernel32Methods, Kernel32Save, Kernel32Count, list);
		ReadHookMethods(user32Methods, user32Save, user32Count, list);
		ReadHookMethods(shell32Methods, shell32Save, shell32Count, list);
		ReadHookMethods(wmvcoreMethods, wmvcoreSave, wmvcoreCount, list);
		ReadHookMethods(msvcrtMethods, msvcrtSave, msvcrtCount, list);
		ReadHookMethods(propsysMethods, propsysSave, propsysCount, list);
		ReadHookMethods(ole32Methods, ole32Save, ole32Count, list);
		ReadHookMethods(mfMethods, mfSave, mfCount, list);
		ReadHookMethods(quartzExternals, quartzSave, classObjectCount, list);
		while(list->GetCount()) {
			LPCSTR name = list->Get(0);
			PHookList list2 = list->GetData(0);
			if (list2) {
				bool notsaved = true;
				while(list2->GetCount()) {
					HookData* hd = &list2->GetAt(0);
					if (notsaved) {
						SaveHookMethod(name, hd->hook);
						notsaved = false;
					}
					INT_PTR m = RestoreModules->IndexOf(hd->module);
					if (m != -1) {
						PSetPrevAddress restore = RestoreModules->GetData(m);
						if (list2->GetCount() > 1) {
							HookData* hd2 = &list2->GetAt(1);
							(*restore)(name, hd2->hook);
						} else if (hd->save) {
							(*restore)(name, hd->save);
						} else {
							RestoreProcs->Add(name, restore);
						}
					}
					list2->RemoveAt(0);
				}
			}
			list->Remove(0);
		}
	} finally (
		delete list;
	)
}

void Uninitialize() {
	HookUninitializing = true;
	bool exit;
	monitor->AcquireLockShared();
	exit = Initialized;
	monitor->ReleaseLockShared();
	if (exit)
		return;
	Locker->AcquireLockExclusive();
	try {
		CSLocker->AcquireLockExclusive();
		try {
			SaveHookMethod(NULL, NULL);
			RestoreProcs->Clear();
			while(RestoreModules->GetCount()) {
				HMODULE module = RestoreModules->Get(0);
				try {
					void* uninit = GetProcAddress(module, "WMPUninitialize");
					if (uninit) {
						(*(PWMPUninitialize)uninit)();
					}
				} finally (
					FreeLibrary(module);
				)
				RestoreModules->Remove(0);
			}
		} catch(...) {
			TraceException(_T("Uninitialize"));
		}
		CSLocker->ReleaseLockExclusive();
	} finally (
		Locker->ReleaseLockExclusive();
	)
}

void Initialize(bool nolock) {
	CreatePathBuffer(pathBuffer);
	CreatePathBuffer(nameBuffer);

	if (Initialized)
		return;
	monitor->AcquireLockExclusive();
	if (!Initialized) {
		Initialized = true;
		try {
			if (!RestoreModules) {
				RestoreModules = new RestoreModuleList();
			}
			if (!RestoreProcs) {
				RestoreProcs = new RestoreProcList();
			}
			if (!Locker) {
				Locker = CreateRWLocker();
			}
			if (!nolock) {
				if (!LockModuleIntoProcess(HInstance))
					OutputDebugString(_T("LockModuleIntoProcess failed\r\n"));
			}
			GetModuleFileName(NULL, pathBuffer, pathBufferLength);
			bool iswmplayer = lstrcmpi(GetFileName(pathBuffer, nameBuffer), WMPLAYEREXE) == 0;
			bool hook = iswmplayer && !CommandParamExist(GetCommandLine(), _T("/nowmphook"));
			int patched = 0;
			if (hook) {
				HMODULE handle = GetModuleHandle(SHELL32DLL);
				if (!handle)
					handle = LoadLibrary(SHELL32DLL);
				patched |= Patch(handle, SHELL32DLL, PATCH_ACTION_SHELL32);
				handle = GetModuleHandle(OLE32DLL);
				if (!handle)
					handle = LoadLibrary(OLE32DLL);
				patched |= Patch(handle, OLE32DLL, PATCH_ACTION_OLE32);
			}
			GetSystemDirectory(pathBuffer, pathBufferLength);
			FixPath(pathBuffer);
			_tcscat(pathBuffer, _T("\\wmp.dll"));
			HINST = LoadLibrary(pathBuffer);
			if (HINST) {
				patched |= Patch(HINST, pathBuffer,  PATCH_ACTION_WMP);
				ReadMethods(wmpSave, HINST, wmpMethods, wmpOrdinals, wmpCount);
				if (iswmplayer && hook) {
					InitializeModule(HINST, (IMPORT_LIBRARY)patched);
				}
			} else {
				OutputDebugString(_T("Error: failed to load wmp.dll\r\n"));
			}
		} catch(...) {
			TraceException(_T("Initialize"));
		}
	}
	monitor->ReleaseLockExclusive();
}

HRESULT WINAPI StartThePlayerProxy(void* arg1, INT_PTR arg2, int arg3) {
	Initialize();
	void* save = wmpSave[0];
	if (save) {
		return (*(PStartThePlayer)save)(arg1, arg2, arg3);
	}
	return S_OK;
}

HRESULT WINAPI CheckThePlayerProxy(LPSTR* name, PDWORD majorFileVer, PDWORD minorFileVer, PDWORD majorRegVer, PDWORD minorRegVer) {
	Initialize();
	void* save = wmpSave[1];
	if (save) {
		return (*(PCheckThePlayer)save)(name, majorFileVer, minorFileVer, majorRegVer, minorRegVer);
	}
	return S_OK;
}

void WINAPI ReleaseDataObjProxy(LPCWSTR path) {
	Initialize();
	void* save = wmpSave[2];
	if (save) {
		(*(PReleaseDataObj)save)(path);
	}
}

HRESULT WINAPI UpdateSharedFoldersProxy() {
	Initialize();
	void* save = wmpSave[3];
	if (save) {
		return (*(PUpdateSharedFolders)save)();
	}
	return S_OK;
}

HRESULT CallFunctionElevatedProxy(LPCWSTR name) {
	Initialize();
	void* save = wmpSave[4];
	if (save) {
		return (*(PCallFunctionElevated)save)(name);
	}
	return S_OK;
}

bool WINAPI ShouldActiveSetupBeRunProxy() {
	Initialize();
	void* save = wmpSave[5];
	if (save) {
		return (*(PShouldActiveSetupBeRun)save)();
	}
	return false;
}

HRESULT WINAPI CheckThePlayer11Proxy(LPSTR* name, PDWORD majorFileVer, PDWORD minorFileVer, PDWORD majorRegVer, PDWORD minorRegVer) {
	Initialize();
	void* save = wmpSave[6];
	if (save) {
		return (*(PCheckThePlayer)save)(name, majorFileVer, minorFileVer, majorRegVer, minorRegVer);
	}
	return S_OK;
}

HRESULT WINAPI DllCanUnloadNow() {
	Initialize();
	void* save = wmpSave[7];
	if (save) {
		return (*(PDllCanUnloadNow)save)();
	}
	return S_OK;
}

HRESULT WINAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID *ppv) {
	Initialize();
	void* save = wmpSave[8];
	if (save) {
		return (*(PDllGetClassObject)save)(rclsid, riid, ppv);
	}
	return E_UNEXPECTED;
}

HRESULT WINAPI DllRegisterServer() {
	Initialize();
	void* save = wmpSave[9];
	if (save) {
		return (*(PDllRegisterServer)save)();
	}
	return E_UNEXPECTED;
}

HRESULT WINAPI DllUnregisterServer() {
	Initialize();
	void* save = wmpSave[10];
	if (save) {
		return (*(PDllUnregisterServer)save)();
	}
	return E_UNEXPECTED;
}

HRESULT WINAPI InitializeTests() {
	Initialize();
	void* save = wmpSave[11];
	if (save) {
		return (*(PInitializeTests)save)();
	}
	return S_OK;
}

void HookHandle(HMODULE* hModule) {
	if (*hModule == HInstance)
		*hModule = HINST;
}

void Cleanup() {
	try {
		if (RestoreModules) {
			delete RestoreModules;
			RestoreModules = NULL;
		}
		if (RestoreProcs) {
			delete RestoreProcs;
			RestoreProcs = NULL;
		}
	} catch(...) {
		TraceException(_T("Cleanup"));
	}
}

void AcquireLock() {
	DWORD lastError = GetLastError();
	try {
		Locker->AcquireLockShared();
	} finally (
		SetLastError(lastError);
	)
}

void ReleaseLock() {
	DWORD lastError = GetLastError();
	try {
		Locker->ReleaseLockShared();
	} finally (
		SetLastError(lastError);
	)
}

WMPHookUnit::WMPHookUnit() {
	monitor = new CriticalSectionLocker();
	CSLocker = new CriticalSectionLocker();
}

WMPHookUnit::~WMPHookUnit() {
	Uninitialized = true;
	delete monitor;
	Cleanup();
	delete CSLocker;
}
